package graph

import (
	"math"
)

// Edge 图中的路径和权重。
type Edge struct {
	To     int
	Weight float64
}

// FloydWarshall Floyd-Warshall算法是解决任意两点间的最短路径的一种算法，可以正确处理有向图或负权的最短路径问题。
// Floyd-Warshall算法的时间复杂度为O(N3)，空间复杂度为O(N2)。
func FloydWarshall(edges [][]Edge) ([][]float64, [][]int) {
	distances := make([][]float64, len(edges))
	nexts := make([][]int, len(edges))
	for i := range distances {
		drow := make([]float64, len(edges))
		nrow := make([]int, len(edges))
		for j := range drow {
			drow[j] = math.Inf(1)
			nrow[j] = -1
		}
		drow[i] = 0
		distances[i] = drow

		nrow[i] = -1
		nexts[i] = nrow
	}

	for i, di := range edges {
		for _, edge := range di {
			distances[i][edge.To] = edge.Weight
			nexts[i][edge.To] = edge.To
		}
	}

	for k, dk := range distances {
		for i, di := range distances {
			for j, dij := range di {
				d := di[k] + dk[j]
				if dij > d {
					di[j] = d
					nexts[i][j] = nexts[i][k]
				}
			}
		}
	}

	return distances, nexts
}

// GetPath 获取两点之间的最短路径。
func getPath(u, v int, nexts [][]int) []int {
	var path []int
	if u != v && nexts[u][v] == -1 {
		return path
	}

	path = append(path, u)
	for u != v {
		u = nexts[u][v]
		if u == -1 {
			return []int{}
		}
		path = append(path, u)
	}

	return path
}

// Graph 图结构.
type Graph struct {
	Nodes map[string]int // node name to node id
	IDs   map[int]string // node id to node name
	// Edges     map[int][]Edge // edges, key is node id
	Edges     [][]Edge    // edges, key is node id
	Distances [][]float64 // shortest distance between two nodes, key is node id i,j
	Nexts     [][]int     // next node id of the shortest path between two nodes, key is node id i,j
}

// NewGraph 创建图。
func NewGraph() *Graph {
	return &Graph{
		Nodes: make(map[string]int),
		IDs:   make(map[int]string),
		// Edges: make(map[int][]Edge),
	}
}

// AddEdge 添加边。
func (g *Graph) AddEdge(u, v string, weight float64, directed bool) {
	i, existed := g.Nodes[u]
	if !existed {
		i = len(g.Nodes)
		g.Nodes[u] = i
		g.IDs[i] = u
		g.Edges = append(g.Edges, []Edge{})
	}
	j, existed := g.Nodes[v]
	if !existed {
		j = len(g.Nodes)
		g.Nodes[v] = j
		g.IDs[j] = v
		g.Edges = append(g.Edges, []Edge{})
	}
	g.Edges[i] = append(g.Edges[i], Edge{To: j, Weight: weight})
	if !directed {
		g.Edges[j] = append(g.Edges[j], Edge{To: i, Weight: weight})
	}
}

// AddDirectedEdge 添加有向边。
func (g *Graph) AddDirectedEdge(u, v string, weight float64) {
	g.AddEdge(u, v, weight, true)
}

// AddUndirectedEdge 添加无向边。
func (g *Graph) AddUndirectedEdge(u, v string, weight float64) {
	g.AddEdge(u, v, weight, false)
}

// HasEdge 检查两点之间是否存在边。
func (g *Graph) HasEdge(u, v string) bool {
	i := g.Nodes[u]
	j := g.Nodes[v]
	for _, edge := range g.Edges[i] {
		if edge.To == j {
			return true
		}
	}
	return false
}

// CalcPaths 计算最短路径。
func (g *Graph) CalcPaths() {
	results, nexts := FloydWarshall(g.Edges)
	g.Distances = results
	g.Nexts = nexts
}

// FindPath 查找两点之间的最短路径。
func (g *Graph) FindPath(u, v string) []string {
	i := g.Nodes[u]
	j := g.Nodes[v]

	paths := getPath(i, j, g.Nexts)

	var namePaths []string
	for _, k := range paths {
		namePaths = append(namePaths, g.IDs[k])
	}

	return namePaths
}

// NeExampleGraph 一个例子
func NeExampleGraph() *Graph {
	g := NewGraph()
	ConstructExampleGraph(g)
	g.CalcPaths()
	return g
}

// ConstructExampleGraph 构造一个例子。
func ConstructExampleGraph(g *Graph) {
	g.AddUndirectedEdge("A", "B", 100)
	g.AddUndirectedEdge("B", "D", 90)
	g.AddUndirectedEdge("A", "C", 100)
	g.AddUndirectedEdge("C", "D", 100)
	g.AddUndirectedEdge("B", "C", 100)
}
